# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：ellipticIntegral.py
    Created on：2025/4/12 22:49
    Author：(input)
    Description:
"""
import numpy as np
from scipy.special import ellipk, ellipj
from scipy.optimize import minimize
import sympy as sp

class EllipticIntegralTransform:
    def __init__(self, k):
        """
        初始化椭圆积分变换，输入模数 k。
        参数：
            k: 模数 (0 <= k <= 1)，可以是 float, int 或 sympy 数值
        """
        # 转换为标准 float，处理 sympy.Float 等类型
        if isinstance(k, (sp.Expr, sp.Float)):
            k = float(k)
        if not isinstance(k, (int, float)):
            raise ValueError(f"k 必须是数值类型，得到 {type(k)}")
        if not (0 <= k <= 1):
            raise ValueError(f"模数 k 必须在 [0, 1]，得到 {k}")
        self.k = k
        self.k_prime = None  # 补模数 k' = sqrt(1 - k^2)
        self.K = None  # 完全椭圆积分 K(k)
        self.K_prime = None  # 完全椭圆积分 K(k')

    def compute_modulus(self):
        """
        计算补模数 k' 和完全椭圆积分 K(k), K(k')。
        """
        self.k_prime = np.sqrt(1 - self.k ** 2)
        self.K = ellipk(self.k ** 2)
        self.K_prime = ellipk(self.k_prime ** 2)
        print(f"模数 k: {self.k}")
        print(f"补模数 k': {self.k_prime}")
        print(f"完全椭圆积分 K(k): {self.K}")
        print(f"完全椭圆积分 K(k'): {self.K_prime}")

    def verify_modulus(self):
        """
        验证 k^2 + k'^2 = 1。
        """
        if self.k_prime is None:
            raise ValueError("请先计算 k'")
        sum_of_squares = self.k ** 2 + self.k_prime ** 2
        if not np.isclose(sum_of_squares, 1, rtol=1e-10):
            raise ValueError(f"验证失败：k^2 + k'^2 = {sum_of_squares}")
        print("验证通过：k^2 + k'^2 = 1")

    def compute_elliptic_integral(self, u):
        """
        计算椭圆积分变换，将实轴坐标 u 映射到复数坐标 z = sn(u, k)。
        参数：
            u: 实数坐标（标量）
        返回：
            z: sn(u, k)，复数值
        """
        if not isinstance(u, (int, float)):
            raise ValueError(f"u 必须是实数标量，得到 {type(u)}")
        try:
            sn, _, _, _ = ellipj(float(u), float(self.k ** 2))
            z = sn  # 返回 sn(u, k)，复数形式
            return z
        except Exception as e:
            raise ValueError(f"计算椭圆积分失败：{e}")

    def verify_elliptic_integral(self, u):
        """
        验证椭圆积分变换的性质，例如 sn(u + 4K, k) = sn(u, k)。
        参数：
            u: 实数坐标（标量）
        """
        if self.K is None:
            raise ValueError("请先计算 K(k)")
        try:
            z = self.compute_elliptic_integral(u)
            # 验证周期性：sn(u + 4K, k) ≈ sn(u, k)
            z_periodic = self.compute_elliptic_integral(u + 4 * self.K)
            print(f"验证椭圆积分：")
            print(f"u = {u}, z = sn(u, k) = {z}")
            print(f"u + 4K = {u + 4 * self.K}, z = sn(u + 4K, k) = {z_periodic}")
            if not np.isclose(z, z_periodic, rtol=1e-5):
                print("警告：周期性验证未通过")
            else:
                print("周期性验证通过！")
            # 验证特殊点：sn(K, k) = 1
            if np.isclose(u, self.K, rtol=1e-5):
                print(f"u ≈ K(k), 检查 sn(K, k) = 1: {z}")
                if not np.isclose(z, 1, rtol=1e-5):
                    print("警告：sn(K, k) ≠ 1")
                else:
                    print("特殊点验证通过！")
            return z, z_periodic
        except Exception as e:
            raise ValueError(f"验证椭圆积分失败：{e}")

    def solve(self):
        """
        执行完整求解流程，初始化模数和验证。
        返回：
            k, k_prime, K, K_prime
        """
        print("初始化模数和补模数...")
        self.compute_modulus()
        print("验证模数...")
        self.verify_modulus()
        return self.k, self.k_prime, self.K, self.K_prime

# 示例用法
if __name__ == "__main__":
    # 定义模数 k
    k = 0.5
    transformer = EllipticIntegralTransform(k)

    # 执行求解流程
    k, k_prime, K, K_prime = transformer.solve()

    # 测试椭圆积分变换
    print("\n测试椭圆积分变换：")
    u_values = [0.0, 1.0, K, 2.0]  # 测试不同的实数输入
    for u in u_values:
        z = transformer.compute_elliptic_integral(u)
        print(f"输入 u = {u:.6f}, 输出 z = sn(u, k) = {z}")

    # 验证椭圆积分性质
    print("\n验证椭圆积分性质：")
    transformer.verify_elliptic_integral(K)  # 测试 u = K
    transformer.verify_elliptic_integral(1.0)  # 测试任意点